#-*- coding:utf-8 -*-

import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data


def add_layer(input, out_size, activation_function=None):
    W = tf.Variable(tf.random_normal([input.shape.as_list()[1],out_size]))
    b = tf.Variable(tf.zeros([out_size])+0.01)

    Z = tf.nn.bias_add(tf.matmul(input,W),b)

    if activation_function is None:
        outputs = Z
    else:
        outputs = activation_function(Z)

    return outputs

if __name__ == '__main__':
    MNIST = input_data.read_data_sets(r"E:\testDir\ml\trainData\mnist",one_hot=True)

    learning_rate = 0.01
    batch_size = 128
    n_epochs = 70

    X = tf.placeholder(tf.float32,shape=[None,784])
    Y = tf.placeholder(tf.float32, shape=[None, 10])

    layer_dims = [784,500,500,10]
    layer_count = len(layer_dims) - 1
    layer_iter = X

    for l in range(1, layer_count):
        layer_iter = add_layer(layer_iter, layer_dims[l], activation_function=tf.nn.relu)
    prediction = add_layer(layer_iter, layer_dims[layer_count], activation_function=None )

    entropy = tf.nn.softmax_cross_entropy_with_logits(labels=Y, logits=prediction)
    loss = tf.reduce_mean(entropy)

    optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)

    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)

        n_batches = int(MNIST.train.num_examples/batch_size)

        for i in range(n_epochs):
             for j in range(n_batches):
                 X_batches,Y_batches = MNIST.train.next_batch(batch_size)
                 _, loss_ = sess.run([optimizer, loss], feed_dict={X:X_batches, Y:Y_batches})
                 if i % 10 == 5 and j == 0:
                     print("Loss of epochs[{0}]: {1}".format(i, loss_))

        # test the model
        n_batches = int(MNIST.test.num_examples/batch_size)
        total_correct_num = 0
        for i in range(n_batches):
            X_batches_test,Y_batches_test = MNIST.test.next_batch(batch_size)
            preds = sess.run(prediction,feed_dict={X:X_batches_test, Y:Y_batches_test})

            correct_preds = tf.equal(tf.argmax(preds,1),tf.argmax(Y_batches_test,1))
            correct_num = tf.reduce_sum(tf.cast(correct_preds,tf.float32))

            total_correct_num += sess.run(correct_num)

        print("Accuracy {0}".format(total_correct_num / MNIST.test.num_examples))